import yara
import os
import sys
 
def getRules(path):
    filepath = {}
    for index,file in enumerate(os.listdir(path)):
        rupath = os.path.join(path, file)
        key = "rule"+str(index)
        filepath[key] = rupath
    yararule = yara.compile(filepaths=filepath)
    return yararule
 
def scan(rule, path):
    count = 0
    for file in os.listdir(path):
        mapath = os.path.join(path, file)
        fp = open(mapath, 'rb')
        matches = rule.match(data=fp.read())
        if len(matches)>0:
            removeFile(malpath,file)
            #print file,matches
            count= count+1

def removeFile(dir,postfix):
    if os.path.isdir(dir):
        for file in os.listdir(dir):
            removeFile(dir+'/'+file,postfix)
    else:
        # if os.path.splitext(dir)[0]==postfix:
        if os.path.basename(dir)==postfix:
            os.remove(dir)

# def logDanger(filename,jsonfile):
    
            
if __name__ == '__main__':
    rulepath = sys.argv[1]
    malpath = sys.argv[2]

    yararule = getRules(r_path)
    scan(yararule, malpath)
    print count